{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "from datasets import load_dataset, Dataset, load_from_disk\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    ")\n",
    "\n",
    "from kw_filter import kw_cls\n",
    "from industry_binary_cls import (\n",
    "    run_llm_binary_cls,\n",
    "    build_binary_dataset,\n",
    "    merge_binary_dataset,\n",
    ")\n",
    "from setting import address, StaticValues\n",
    "from prompt import get_industry_trans_func\n",
    "from bert_train import BertCLS, bert_cls_trans, bert_tokenize_func\n",
    "from utils import tokenize_func, export_human_csv\n",
    "from industry_multi_cls import run_multi_cls, trans_multi_cls_dataset\n",
    "from collections import Counter\n",
    "from datasets import concatenate_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/16 00:57:35 - machine - INFO: \n",
      "开始进入bert训练多分类\n",
      "\n"
     ]
    }
   ],
   "source": [
    "name = \"machine\"\n",
    "sv = StaticValues(name)\n",
    "logger = sv.logger\n",
    "\n",
    "logger.info(\"开始进入bert训练多分类\")\n",
    "\n",
    "bert_multi_dataset = merge_binary_dataset(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(sv.bert_config.bert_multi_train):\n",
    "    logger.info(\"开始使用大模型预测多分类\")\n",
    "    bert_multi_dataset = merge_binary_dataset(name)\n",
    "    bert_multi_dataset = bert_multi_dataset.train_test_split(train_size=30000)\n",
    "    bert_multi_dataset[\"train\"] = run_multi_cls(name, bert_multi_dataset[\"train\"])\n",
    "    bert_multi_dataset[\"train\"] = bert_multi_dataset[\"train\"].map(\n",
    "        trans_multi_cls_dataset(name)\n",
    "    )\n",
    "    bert_multi_dataset.save_to_disk(sv.bert_config.bert_multi_train)\n",
    "else:\n",
    "    bert_multi_dataset = load_from_disk(sv.bert_config.bert_multi_train)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bert_cls_trans(name, bert_multi_dataset):\n",
    "    \"\"\"\n",
    "    把多标签分类的数据集转成多类别分类数据集\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "    from collections import Counter\n",
    "\n",
    "\n",
    "    sv = StaticValues(name=name)\n",
    "    LABEL_NAME = sv.LABEL_NAME\n",
    "    # bert_multi_dataset = load_from_disk(sv.bert_config.bert_multi_train)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(sv.bert_config.model_name)\n",
    "\n",
    "    def bert_trans_train_dataset(item):\n",
    "        if \"industry_info\" not in item.keys():\n",
    "            item = get_industry_trans_func(\"industry_info\", \"{industry_info}\")(item)\n",
    "        label = 0\n",
    "        for k in LABEL_NAME:\n",
    "            label += item[k]\n",
    "        if label == 1:\n",
    "            for idx, k in enumerate(LABEL_NAME):\n",
    "                if item[k] == 1:\n",
    "                    label = idx\n",
    "        else:\n",
    "            label = -1\n",
    "        item[\"label\"] = label\n",
    "\n",
    "        tokenized_inputs = tokenizer(\n",
    "            item[\"industry_info\"],\n",
    "            max_length=512,\n",
    "            truncation=True,\n",
    "        )\n",
    "        return tokenized_inputs\n",
    "\n",
    "    new_dataset = bert_multi_dataset.map(bert_trans_train_dataset)\n",
    "\n",
    "    idx = np.array(new_dataset[\"label\"]) >= 0\n",
    "    data = np.array(new_dataset[\"label\"])[idx]\n",
    "    min_len = min(Counter(data).values())\n",
    "\n",
    "    def cut_dataset(_dataset, min_len):\n",
    "        length = min(max(3000, 2 * min_len), len(_dataset) - 1)\n",
    "        return _dataset.train_test_split(train_size=length)[\"train\"]\n",
    "\n",
    "    dataset_d = {\n",
    "        k: cut_dataset(new_dataset.filter(lambda item: item[\"label\"] == k), min_len)\n",
    "        for k in range(len(LABEL_NAME))\n",
    "    }\n",
    "    \n",
    "    res_dataset = concatenate_datasets(dataset_d.values())\n",
    "    \n",
    "    logger.info(f\"{Counter(res_dataset['label'])}\")\n",
    "    return res_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/16 00:57:56 - machine - INFO: \n",
      "开始进入bert训练多分类\n",
      "\n",
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95d77af126b1447bbf996493bd7537fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4a17fd1d1f745fe8903401729a3eb98",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7212d141c61430ea583ed5f94abcb95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "473d6ab7cd4045aea58631f8961fd2b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "174d14eb9a004c0882077c6c8e69595d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/16 00:58:32 - machine - INFO: \n",
      "Counter({3: 3000, 0: 1859, 2: 1096, 1: 273})\n",
      "\n",
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/training_args.py:1525: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "logger.info(\"开始进入bert训练多分类\")\n",
    "new_dataset = bert_cls_trans(name, bert_multi_dataset[\"train\"])\n",
    "\n",
    "multi_best_model = os.path.join(sv.bert_config.output_multi_dir, \"best_model\")\n",
    "if not os.path.exists(multi_best_model):\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        sv.bert_config.model_name, num_labels=len(sv.LABEL_NAME)\n",
    "    )\n",
    "    tokenizer = AutoTokenizer.from_pretrained(sv.bert_config.model_name)\n",
    "    bert_cls = BertCLS(\n",
    "        model=model,\n",
    "        tokenizer=tokenizer,\n",
    "        train_dataset=new_dataset,\n",
    "        output_dir=sv.bert_config.output_multi_dir,\n",
    "    )\n",
    "    bert_cls.train()\n",
    "else:\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        multi_best_model, num_labels=len(sv.LABEL_NAME)\n",
    "    )\n",
    "    tokenizer = AutoTokenizer.from_pretrained(sv.bert_config.model_name)\n",
    "    bert_cls = BertCLS(model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/16 00:59:13 - machine - INFO: \n",
      "bert 多分类预测开始\n",
      "\n",
      "2024/09/16 00:59:13 - machine - INFO: \n",
      "get_industry_trans_func...\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fda6d9215feb42f08f6e833d979e7bcc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/200 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/16 00:59:13 - machine - INFO: \n",
      "bert_tokenize_func...\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32630ad3e52d49e4bb7b9b4806b5a09d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/200 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-09-16 00:59:13,778] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jie/anaconda3/envs/llm/compiler_compat/ld: cannot find -laio: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n",
      "/home/jie/anaconda3/envs/llm/compiler_compat/ld: cannot find -lcufile: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n",
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mjieshen\u001b[0m (\u001b[33mjieshenai\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.18.0 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.0"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/jie/gitee/pku_industry/general/wandb/run-20240916_005916-qpdmnr3z</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/jieshenai/huggingface/runs/qpdmnr3z' target=\"_blank\">output</a></strong> to <a href='https://wandb.ai/jieshenai/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/jieshenai/huggingface' target=\"_blank\">https://wandb.ai/jieshenai/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/jieshenai/huggingface/runs/qpdmnr3z' target=\"_blank\">https://wandb.ai/jieshenai/huggingface/runs/qpdmnr3z</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67193e1f932244b3ab7422b2aeac57db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating CSV from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a6e39c5e6574875a4a1d346e1f02f6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89f15f610683400ca481265b96224d62",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/200 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jie/gitee/pku_industry/general/utils.py:144: DtypeWarning: Columns (1,10,11,16,18,23,25,26,27,28,31) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "  kw_df = pd.read_csv(sv.KW_CSV)\n"
     ]
    }
   ],
   "source": [
    "logger.info(\"bert 多分类预测开始\")\n",
    "pred_multi_csv = os.path.join(sv.home_folder, \"bert_multi_pred.csv\")\n",
    "if not os.path.exists(pred_multi_csv):\n",
    "    pred_dataset = bert_multi_dataset[\"test\"]\n",
    "    # test, TODO\n",
    "    pred_dataset = pred_dataset.train_test_split(train_size=200)[\"train\"]\n",
    "    logger.info(\"get_industry_trans_func...\")\n",
    "    pred_dataset = pred_dataset.map(\n",
    "        get_industry_trans_func(\"industry_info\", \"{industry_info}\")\n",
    "    )\n",
    "    logger.info(\"bert_tokenize_func...\")\n",
    "    pred_dataset = pred_dataset.map(bert_tokenize_func(tokenizer), batched=True)\n",
    "    pred_ids = bert_cls.predict(pred_dataset)\n",
    "    pred_dataset = pred_dataset.add_column(\"pred_label\", pred_ids)\n",
    "    pred_dataset.to_csv(pred_multi_csv, index=False)\n",
    "\n",
    "export_human_csv_file = os.path.join(sv.home_folder, f\"{sv.chinese_name}.csv\")\n",
    "if not os.path.exists(export_human_csv_file):\n",
    "    export_human_csv(name=name, output_file=export_human_csv_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
